#!/bin/bash
set -euo pipefail

# 默认使用当前环境 Python，可通过 PYTHON 环境变量覆盖
PYTHON_BIN="${PYTHON:-python}"

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
export USE_LWG=False

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${PROJECT_ROOT}"

# ----------- 可配置参数（可用环境变量覆盖）-----------
CONFIG_PATH="${CONFIG_PATH:-${PROJECT_ROOT}/configs/r50_deformable_detr_motip_maritime_640_lwg.yaml}"
CHECKPOINT_PATH="${CHECKPOINT_PATH:-/home/liuyonghui/code/MA-MOTIP/pretrains/maritime_baseline.pth}"
DATA_ROOT="${DATA_ROOT:-/home/liuyonghui/datasets}"
INFERENCE_SPLIT="${INFERENCE_SPLIT:-train}"
MARITIME_DATASET="${MARITIME_DATASET:-MaritimeTrack_Full_20250413}"

LWG_EXPORT_DIR="${LWG_EXPORT_DIR:-${PROJECT_ROOT}/outputs/maritime/lwg_export}"
LWG_DUMP_DIR="${LWG_DUMP_DIR:-${PROJECT_ROOT}/outputs/maritime/lwg_feature_dump}"
LWG_RAW_CSV="${LWG_RAW_CSV:-${PROJECT_ROOT}/outputs/maritime/lwg_training/lwg_train_raw.csv}"
LWG_CSV_PATH="${LWG_CSV_PATH:-${PROJECT_ROOT}/outputs/maritime/lwg_training/lwg_train.csv}"
LWG_MODEL_PATH="${LWG_MODEL_PATH:-${PROJECT_ROOT}/pretrains/lwg_maritime.pt}"

BUFFER_EXPORT_DIR="${BUFFER_EXPORT_DIR:-${PROJECT_ROOT}/outputs/maritime/buffer_gate_export}"
BUFFER_DUMP_DIR="${BUFFER_DUMP_DIR:-${PROJECT_ROOT}/outputs/maritime/buffer_gate_dump}"
BUFFER_RAW_CSV="${BUFFER_RAW_CSV:-${PROJECT_ROOT}/outputs/maritime/buffer_gate_training/buffer_gate_raw.csv}"
BUFFER_MODEL_PATH="${BUFFER_MODEL_PATH:-${PROJECT_ROOT}/pretrains/buffer_gate_maritime.pt}"

# LWG 训练参数
LWG_K="${LWG_K:-20}"
LWG_M="${LWG_M:-8}"
LWG_IOU_THRESH="${LWG_IOU_THRESH:-0.8}"
LWG_EPOCHS="${LWG_EPOCHS:-200}"
LWG_BATCH_SIZE="${LWG_BATCH_SIZE:-1024}"
LWG_LR="${LWG_LR:-0.001}"
LWG_VAL_RATIO="${LWG_VAL_RATIO:-0.1}"
LWG_PATIENCE="${LWG_PATIENCE:-5}"
SEED="${SEED:-42}"

# Buffer Gate 训练参数
BUFFER_EPOCHS="${BUFFER_EPOCHS:-50}"
BUFFER_BATCH_SIZE="${BUFFER_BATCH_SIZE:-2048}"
BUFFER_LR="${BUFFER_LR:-0.001}"
BUFFER_VAL_RATIO="${BUFFER_VAL_RATIO:-0.1}"
BUFFER_PATIENCE="${BUFFER_PATIENCE:-5}"
BUFFER_POS_WEIGHT="${BUFFER_POS_WEIGHT:-1.0}"
BUFFER_HISTORY_WINDOW="${BUFFER_HISTORY_WINDOW:-3}"

mkdir -p "${LWG_EXPORT_DIR}" "${LWG_DUMP_DIR}" "$(dirname "${LWG_CSV_PATH}")" "$(dirname "${LWG_MODEL_PATH}")"
mkdir -p "${BUFFER_EXPORT_DIR}" "${BUFFER_DUMP_DIR}" "$(dirname "${BUFFER_RAW_CSV}")" "$(dirname "${BUFFER_MODEL_PATH}")"

rm -f "${LWG_RAW_CSV}" "${LWG_CSV_PATH}" "${BUFFER_RAW_CSV}"
rm -rf "${LWG_DUMP_DIR:?}"/* "${BUFFER_DUMP_DIR:?}"/*

echo "=============================="
echo "Step 1: 无 LWG 导出训练特征 (${INFERENCE_SPLIT})"
echo "  Config     : ${CONFIG_PATH}"
echo "  Checkpoint : ${CHECKPOINT_PATH}"
echo "  Dump dir   : ${LWG_DUMP_DIR}"
echo "=============================="

"${PYTHON_BIN}" - <<PY
from pathlib import Path
from compare_lwg import load_config, prepare_common
from submit_and_evaluate import submit_and_evaluate

config_path = Path("${CONFIG_PATH}")
checkpoint = Path("${CHECKPOINT_PATH}")
output_root = Path("${LWG_EXPORT_DIR}").resolve()
data_root = "${DATA_ROOT}"
dump_dir = Path("${LWG_DUMP_DIR}").resolve()

cfg = load_config(str(config_path))
base_cfg = prepare_common(cfg, str(checkpoint), output_root, data_root)
base_cfg["USE_LWG"] = False
base_cfg["USE_BUFFER_GATE"] = False
base_cfg["EXPORT_LWG_FEATURES"] = True
base_cfg["LWG_FEATURE_DUMP_DIR"] = str(dump_dir)
base_cfg["EXPORT_BUFFER_GATE_FEATURES"] = False
base_cfg["INFERENCE_SPLIT"] = "${INFERENCE_SPLIT}"
base_cfg["EXP_NAME"] = base_cfg.get("EXP_NAME", "maritime_lwg_export")
base_cfg["INFERENCE_GROUP"] = base_cfg.get("INFERENCE_GROUP", "maritime_lwg_export")

submit_and_evaluate(base_cfg)
PY

echo "=============================="
echo "Step 2: 构建 LWG 训练 CSV"
echo "  Dump dir : ${LWG_DUMP_DIR}"
echo "  CSV path : ${LWG_RAW_CSV}"
echo "=============================="

build_cmd=(
  tools/lwg/build_dataset.py
  --dump-dir "${LWG_DUMP_DIR}"
  --data-root "${DATA_ROOT}"
  --split "${INFERENCE_SPLIT}"
  --out-path "${LWG_RAW_CSV}"
  --K "${LWG_K}"
  --M "${LWG_M}"
  --iou-thresh "${LWG_IOU_THRESH}"
  --dataset "${MARITIME_DATASET}"
)
"${PYTHON_BIN}" "${build_cmd[@]}"

echo "=============================="
echo "Step 3: 清洗 LWG 数据"
echo "  Raw CSV  : ${LWG_RAW_CSV}"
echo "  Clean CSV: ${LWG_CSV_PATH}"
echo "=============================="

clean_cmd=(
  tools/lwg/clean_dataset.py
  --input "${LWG_RAW_CSV}"
  --output "${LWG_CSV_PATH}"
  --min-area-ratio 5e-5
  --max-area-ratio 0.98
  --min-aspect-ratio 0.05
  --max-aspect-ratio 20.0
  --max-neg-pos-ratio 4.0
  --seed "${SEED}"
)
"${PYTHON_BIN}" "${clean_cmd[@]}"

echo "=============================="
echo "Step 4: 训练 LWG 模型"
echo "  CSV   : ${LWG_CSV_PATH}"
echo "  Model : ${LWG_MODEL_PATH}"
echo "=============================="

if [ -f "${LWG_MODEL_PATH}" ]; then
  BACKUP_PATH="${LWG_MODEL_PATH%.*}_backup_$(date +%Y%m%d_%H%M%S).${LWG_MODEL_PATH##*.}"
  cp "${LWG_MODEL_PATH}" "${BACKUP_PATH}"
  echo "备份原 LWG 模型 -> ${BACKUP_PATH}"
fi

lwg_cmd=(
  tools/train_lwg.py
  --csv "${LWG_CSV_PATH}"
  --out "${LWG_MODEL_PATH}"
  --epochs "${LWG_EPOCHS}"
  --batch-size "${LWG_BATCH_SIZE}"
  --lr "${LWG_LR}"
  --device "cuda"
  --val-ratio "${LWG_VAL_RATIO}"
  --patience "${LWG_PATIENCE}"
  --seed "${SEED}"
)
"${PYTHON_BIN}" "${lwg_cmd[@]}"

echo "=============================="
echo "Step 5: 使用 LWG 导出 Buffer Gate 样本 (${INFERENCE_SPLIT})"
echo "  Dump dir : ${BUFFER_DUMP_DIR}"
echo "=============================="

"${PYTHON_BIN}" - <<PY
from pathlib import Path
from compare_lwg import load_config, prepare_common
from submit_and_evaluate import submit_and_evaluate

config_path = Path("${CONFIG_PATH}")
checkpoint = Path("${CHECKPOINT_PATH}")
output_root = Path("${BUFFER_EXPORT_DIR}").resolve()
data_root = "${DATA_ROOT}"
dump_dir = Path("${BUFFER_DUMP_DIR}").resolve()

cfg = load_config(str(config_path))
base_cfg = prepare_common(cfg, str(checkpoint), output_root, data_root)
base_cfg["USE_LWG"] = True
base_cfg["LWG_MODEL_PATH"] = str(Path("${LWG_MODEL_PATH}").resolve())
base_cfg["USE_BUFFER_GATE"] = False
base_cfg["EXPORT_LWG_FEATURES"] = False
base_cfg["EXPORT_BUFFER_GATE_FEATURES"] = True
base_cfg["BUFFER_GATE_DUMP_DIR"] = str(dump_dir)
base_cfg["BUFFER_HISTORY_WINDOW"] = ${BUFFER_HISTORY_WINDOW}
base_cfg["INFERENCE_SPLIT"] = "${INFERENCE_SPLIT}"
base_cfg["EXP_NAME"] = base_cfg.get("EXP_NAME", "maritime_buffer_gate_export")
base_cfg["INFERENCE_GROUP"] = base_cfg.get("INFERENCE_GROUP", "maritime_buffer_gate_export")

submit_and_evaluate(base_cfg)
PY

echo "=============================="
echo "Step 6: 汇总 Buffer Gate 样本"
echo "  Dump dir : ${BUFFER_DUMP_DIR}"
echo "  Raw CSV : ${BUFFER_RAW_CSV}"
echo "=============================="

"${PYTHON_BIN}" - <<PY
from pathlib import Path
import csv

dump_dir = Path("${BUFFER_DUMP_DIR}")
out_path = Path("${BUFFER_RAW_CSV}")
rows = []
header = None
if dump_dir.exists():
    for file in sorted(dump_dir.glob("*.csv")):
        with file.open("r") as f:
            reader = csv.reader(f)
            local_header = next(reader, None)
            if local_header is None:
                continue
            if header is None:
                header = local_header
            for row in reader:
                rows.append(row)
if header and rows:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(rows)
    print(f"[INFO] Collected {len(rows)} buffer samples -> {out_path}")
else:
    print("[WARN] No buffer gate samples collected.")
PY

if [ -s "${BUFFER_RAW_CSV}" ]; then
  echo "=============================="
  echo "Step 7: 训练 Buffer Gate 模型"
  echo "  CSV   : ${BUFFER_RAW_CSV}"
  echo "  Model : ${BUFFER_MODEL_PATH}"
  echo "=============================="

  if [ -f "${BUFFER_MODEL_PATH}" ]; then
    BG_BACKUP="${BUFFER_MODEL_PATH%.*}_backup_$(date +%Y%m%d_%H%M%S).${BUFFER_MODEL_PATH##*.}"
    cp "${BUFFER_MODEL_PATH}" "${BG_BACKUP}"
    echo "备份原 Buffer Gate 模型 -> ${BG_BACKUP}"
  fi

  buffer_cmd=(
    tools/train_buffer_gate.py
    --csv "${BUFFER_RAW_CSV}"
    --out "${BUFFER_MODEL_PATH}"
    --epochs "${BUFFER_EPOCHS}"
    --batch-size "${BUFFER_BATCH_SIZE}"
    --lr "${BUFFER_LR}"
    --device "cuda"
    --val-ratio "${BUFFER_VAL_RATIO}"
    --patience "${BUFFER_PATIENCE}"
    --seed "${SEED}"
    --pos-weight "${BUFFER_POS_WEIGHT}"
  )
  "${PYTHON_BIN}" "${buffer_cmd[@]}"
else
  echo "[WARN] Buffer gate 数据为空，跳过训练。"
fi

echo "=============================="
echo "完成！"
echo "  LWG 模型        : ${LWG_MODEL_PATH}"
echo "  Buffer Gate 模型: ${BUFFER_MODEL_PATH}"
echo "=============================="
